module operators_adv_mod
  use gptl_timer_mod
  use spark_cubed_sphere
  use params_mod
  use static_mod
  use state_mod
  use diag_mod
  use tend_mod
  use math_mod
  use time_mod
  use riemann_mod
  use ghost_interp_mod
  use hope_recon_mod
  use test_mod

  implicit none

  private

  public operators_adv_run

contains

  subroutine operators_adv_run(state, tend, substep)
    type(state_type), intent(inout) :: state
    type(tend_type ), intent(inout) :: tend
    integer, intent(in) :: substep
    
    if(substep==1) call test_adv_update_wind(case_id, elapsed_seconds, diag)
    
    !ret = gptlstart('add_ghs')
    !call add_ghs(state, diag)
    !ret = gptlstop('add_ghs')

    ret = gptlstart('reconstruct')
    call reconstruct(state, diag)
    ret = gptlstop('reconstruct')

    !ret = gptlstart('remove_ghs')
    !call remove_ghs(state, diag)
    !ret = gptlstop('remove_ghs')

    !ret = gptlstart('calc_grad_gh')
    !call calc_grad_gh(state, diag)
    !ret = gptlstop('calc_grad_gh')

    ret = gptlstart('unify_domain_edges')
    call unify_domain_edges(state, diag)
    ret = gptlstop('unify_domain_edges')

    ret = gptlstart('calc_numerical_flux')
    call calc_numerical_flux(state, diag)
    ret = gptlstop('calc_numerical_flux')

    !ret = gptlstart('calc_source_terms')
    !call calc_source_terms(state, diag)
    !ret = gptlstop('calc_source_terms')

    ret = gptlstart('calc_tendency')
    call calc_tendency(state, diag, tend)
    ret = gptlstop('calc_tendency')

  end subroutine operators_adv_run

  subroutine unify_domain_edges(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag

    integer idom, idom_ngb, orient, orient_adj
    integer i, j, k, p, i_ngb, j_ngb, p_ngb, i_adj, j_adj, p_adj
    integer ids, ide, jds, jde, kds, kde, pes(6), pee(6)
    integer is, ie, js, je, di, dj
    real(r8), pointer, dimension(:,:,:,:,:) :: q, q_ngb
    real(r8) :: Jh, Jh_ngb
    real(r8) gd, uc, vc, us, vs

    associate (mesh => state%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde, pes=pes, pee=pee)
        do orient = 1, 4
          if (mesh%domain(idom)%touch_edges(orient)) then
            idom_ngb = mesh%domain(idom)%ngb(orient)%ptr%id
            associate ( ngb => mesh%domain(idom)%ngb(orient) )
            is = ids; ie = ide; js = jds; je = jde; di = 0; dj = 0
            if (orient == 1) then; ie = ids; di = -1; orient_adj = 2; q => diag%domain(idom)%qR; end if
            if (orient == 2) then; is = ide; di = +1; orient_adj = 1; q => diag%domain(idom)%qL; end if
            if (orient == 3) then; je = jds; dj = -1; orient_adj = 4; q => diag%domain(idom)%qT; end if
            if (orient == 4) then; js = jde; dj = +1; orient_adj = 3; q => diag%domain(idom)%qB; end if
            if (ngb%orient == 1) q_ngb => diag%domain(idom_ngb)%qL
            if (ngb%orient == 2) q_ngb => diag%domain(idom_ngb)%qR
            if (ngb%orient == 3) q_ngb => diag%domain(idom_ngb)%qB
            if (ngb%orient == 4) q_ngb => diag%domain(idom_ngb)%qT
            !$OMP PARALLEL DO PRIVATE(j, i, p, i_ngb, j_ngb, p_ngb, i_adj, j_adj, p_adj, Jh, Jh_ngb, gd) COLLAPSE(4)
            do k = kds, kde
              do j = js, je
                do i = is, ie
                  do p = pes(orient), pee(orient)
                    p_adj = p - pes(orient) + pes(orient_adj)
                    i_adj = i + di
                    j_adj = j + dj
                    call mesh%translate_idx_along_edge(idom, orient, i, j, p, i_ngb, j_ngb, p_ngb)
                    Jh     = mesh%get_dom_Jh(idom    ,p_adj,i_adj,j_adj,k)
                    Jh_ngb = mesh%get_dom_Jh(idom_ngb,p_ngb,i_ngb,j_ngb,k)
                    gd = q_ngb(p_ngb,i_ngb,j_ngb,k,1) / Jh_ngb
                    q(p_adj,i_adj,j_adj,k,1) = Jh * gd
                  end do
                end do
              end do
            end do
            !$OMP END PARALLEL DO
            end associate
          end if
        end do
      end if
    end do
    end associate

  end subroutine unify_domain_edges

  subroutine calc_numerical_flux(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag

    integer idom, ids, ide, jds, jde, kds, kde, pes(6), pee(6), pqs
    integer i, j, k, p, p_adj, ivar

    associate (mesh => state%mesh)
    call mesh%get_params(pes=pes, pee=pee, pqs=pqs)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
        associate (iGh  => mesh  %domain(idom)%iGh,  &
                   ghsL => static%domain(idom)%ghsL, &
                   ghsR => static%domain(idom)%ghsR, &
                   ghsB => static%domain(idom)%ghsB, &
                   ghsT => static%domain(idom)%ghsT, &
                   ucL  => diag %domain(idom)%ucL,   &
                   ucR  => diag %domain(idom)%ucR,   &
                   ucB  => diag %domain(idom)%ucB,   &
                   ucT  => diag %domain(idom)%ucT,   &
                   vcL  => diag %domain(idom)%vcL,   &
                   vcR  => diag %domain(idom)%vcR,   &
                   vcB  => diag %domain(idom)%vcB,   &
                   vcT  => diag %domain(idom)%vcT,   &
                   qL   => diag %domain(idom)%qL,    &
                   qR   => diag %domain(idom)%qR,    &
                   qB   => diag %domain(idom)%qB,    &
                   qT   => diag %domain(idom)%qT,    &
                   fxp  => diag %domain(idom)%fxp,   &
                   fx   => diag%domain(idom)%fx,     &
                   fyp  => diag %domain(idom)%fyp,   &
                   fy   => diag%domain(idom)%fy      )
        !$OMP PARALLEL
        !$OMP DO PRIVATE(p, p_adj, ivar) COLLAPSE(3)
        do k = kds, kde
          do j = jds, jde
            do i = ids, ide + 1
              do p = pes(1), pee(1)
                p_adj = p - pes(1) + pes(2)
                fxp(p,i,j,k,:) = riemann_solver_adv(ucL(p,i,j,k), qR(p_adj,i-1,j,k,:), qL(p,i,j,k,:), ghsR(p_adj,i-1,j))
              end do
              do ivar = 1, nvar
                fx(i,j,k,ivar) = mesh%domain(idom)%edge_quad(fxp(:,i,j,k,ivar))
              end do
            end do
          end do
        end do
        !$OMP END DO NOWAIT
        !$OMP DO PRIVATE(p, p_adj, ivar) COLLAPSE(3)
        do k = kds, kde
          do j = jds, jde + 1
            do i = ids, ide
              do p = pes(3), pee(3)
                p_adj = p - pes(3) + pes(4)
                fyp(p,i,j,k,:) = riemann_solver_adv(vcB(p,i,j,k), qT(p_adj,i,j-1,k,:), qB(p,i,j,k,:), ghsT(p_adj,i,j-1))
              end do
              do ivar = 1, nvar
                fy(i,j,k,ivar) = mesh%domain(idom)%edge_quad(fyp(:,i,j,k,ivar))
              end do
            end do
          end do
        end do
        !$OMP END DO
        !$OMP END PARALLEL
        end associate
      end if
    end do
    end associate

  end subroutine calc_numerical_flux

  subroutine calc_source_terms(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    
    real(r8), dimension(ncq,3) :: psi_M, psi_C, psi_B
    real(r8), dimension(ncq) :: gd, u, v, gdu, gdv, gduu, gduv, gdvv, Jhddel2, tanx2, tany2, Jhghs
    integer idom, ids, ide, jds, jde, kds, kde, pc, pqs, pqe, i, j, k, ivar
    
    associate (mesh => state%mesh)
    call mesh%get_params(kds=kds, kde=kde, pc=pc, pqs=pqs, pqe=pqe)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(ids=ids, ide=ide, jds=jds, jde=jde)
        associate ( src => diag%domain(idom)%src )
        !$OMP PARALLEL DO COLLAPSE(3)
        do k = kds, kde
          do j = jds, jde
            do i = ids, ide
              src(i,j,k,:) = 0
            end do
          end do
        end do
        !$OMP END PARALLEL DO
        end associate
      end if
    end do
    end associate

  end subroutine calc_source_terms
  
  subroutine calc_tendency(state, diag, tend)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    type(tend_type ), intent(inout) :: tend

    integer idom, ids, ide, jds, jde, kds, kde, i, j, k, ivar

    associate (mesh => state%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
        associate(dqdt => tend%domain(idom)%q,   &
                  src  => diag%domain(idom)%src, &
                  fx   => diag%domain(idom)%fx,  &
                  fy   => diag%domain(idom)%fy,  &
                  dx   => mesh%domain(idom)%dx,  &
                  dy   => mesh%domain(idom)%dy   )
        !$OMP PARALLEL DO COLLAPSE(4)
        do ivar = 1, nvar
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                !dqdt(i,j,k,ivar) = - (fx(i+1,j,k,ivar) - fx(i,j,k,ivar)) / dx &
                !                   - (fy(i,j+1,k,ivar) - fy(i,j,k,ivar)) / dy &
                !                   + src(i,j,k,ivar)
                dqdt(i,j,k,ivar) = - (fx(i+1,j,k,ivar) - fx(i,j,k,ivar)) / dx &
                                   - (fy(i,j+1,k,ivar) - fy(i,j,k,ivar)) / dy
              end do
            end do
          end do
        end do
        !$OMP END PARALLEL DO
        end associate
      end if
    end do
    end associate

  end subroutine calc_tendency

  subroutine add_ghs(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag

    real(r8) :: Jh
    integer idom, ims, ime, jms, jme, kds, kde, pc, i, j, k

    associate (mesh => state%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%initialized) then
        call mesh%get_params(idom=idom, ims=ims, ime=ime, jms=jms, jme=jme, kds=kds, kde=kde, pc=pc)
        associate (q   => state%domain(idom)%q, &
                   ghs => static%domain(idom)%ghs)
        !$OMP PARALLEL DO PRIVATE(Jh) COLLAPSE(3)
        do k = kds, kde
          do j = jms, jme
            do i = ims, ime
              Jh = mesh%get_dom_Jh(idom,pc,i,j,k)
              q(i,j,k,1) = q(i,j,k,1) + Jh * ghs(i,j)
            end do
          end do
        end do
        !$OMP END PARALLEL DO
        end associate
      end if
    end do
    end associate

  end subroutine add_ghs

  subroutine remove_ghs(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag

    real(r8) :: Jh
    integer idom, ims, ime, jms, jme, kds, kde
    integer pc, pes(6), pee(6), pqs, pqe
    integer ip, p, i, j, k

    associate (mesh => state%mesh)
    call mesh%get_params(kds=kds, kde=kde, pc=pc, pes=pes, pee=pee, pqs=pqs, pqe=pqe)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%initialized) then
        call mesh%get_params(idom=idom, ims=ims, ime=ime, jms=jms, jme=jme)
        associate (domain => mesh%domain(idom),        &
                   q      => state%domain(idom)%q,     &
                   qL     => diag%domain(idom)%qL,     &
                   qR     => diag%domain(idom)%qR,     &
                   qB     => diag%domain(idom)%qB,     &
                   qT     => diag%domain(idom)%qT,     &
                   ghs    => static%domain(idom)%ghs , &
                   ghsL   => static%domain(idom)%ghsL, &
                   ghsR   => static%domain(idom)%ghsR, &
                   ghsB   => static%domain(idom)%ghsB, &
                   ghsT   => static%domain(idom)%ghsT  )
        !$OMP PARALLEL DO PRIVATE(ip, p, Jh) COLLAPSE(3)
        do k = kds, kde
          do j = jms, jme
            do i = ims, ime
              Jh = mesh%get_dom_Jh(idom,pc,i,j,k)
              q(i,j,k,1) = q(i,j,k,1) - Jh * ghs (i,j)
              
              do ip = 1, neq
                p = pes(1) + ip - 1
                Jh = mesh%get_dom_Jh(idom,p,i,j,k)
                qL(:,i,j,k,1) = qL(:,i,j,k,1) - Jh * ghsL(:,i,j)
                
                p = pes(2) + ip - 1
                Jh = mesh%get_dom_Jh(idom,p,i,j,k)
                qR(:,i,j,k,1) = qR(:,i,j,k,1) - Jh * ghsR(:,i,j)
                
                p = pes(3) + ip - 1
                Jh = mesh%get_dom_Jh(idom,p,i,j,k)
                qB(:,i,j,k,1) = qB(:,i,j,k,1) - Jh * ghsB(:,i,j)
                
                p = pes(4) + ip - 1
                Jh = mesh%get_dom_Jh(idom,p,i,j,k)
                qT(:,i,j,k,1) = qT(:,i,j,k,1) - Jh * ghsT(:,i,j)
              enddo
            end do
          end do
        end do
        !$OMP END PARALLEL DO
        end associate
      end if
    end do
    end associate

  end subroutine remove_ghs

  subroutine calc_grad_gh(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag

    real(r8) :: Jh
    integer idom, ids, ide, jds, jde, kds, kde, pc, i, j, k
    integer       irs, ire, jrs, jre

    associate (mesh => state%mesh)
    call mesh%get_params(kds=kds, kde=kde, pc=pc)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, &
                                        irs=irs, ire=ire, jrs=jrs, jre=jre)
        associate (dx    => mesh%domain(idom)%dx,    &
                   dy    => mesh%domain(idom)%dy,    &
                   ghs   => static%domain(idom)%ghs, &
                   q     => state%domain(idom)%q,    &
                   ght   => diag%domain(idom)%ght,   &
                   dghdy => diag%domain(idom)%dghdy, &
                   dghdx => diag%domain(idom)%dghdx  )
        !$OMP PARALLEL DO PRIVATE(Jh) COLLAPSE(3)
        do k = kds, kde
          do j = jrs, jre
            do i = irs, ire
              Jh = mesh%get_dom_Jh(idom,pc,i,j,k)
              ght(i,j,k) = q(i,j,k,1) / Jh + ghs(i,j)
            end do
          end do
        end do
        !$OMP END PARALLEL DO
        select case (recon_h_order)
        case (3)
          !$OMP PARALLEL DO COLLAPSE(3)
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                dghdx(i,j,k) = center_diff2(ght(i-1:i+1,j,k), dx)
                dghdy(i,j,k) = center_diff2(ght(i,j-1:j+1,k), dy)
              end do
            end do
          end do
          !$OMP END PARALLEL DO
        case (5)
          !$OMP PARALLEL DO COLLAPSE(3)
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                dghdx(i,j,k) = center_diff4(ght(i-2:i+2,j,k), dx)
                dghdy(i,j,k) = center_diff4(ght(i,j-2:j+2,k), dy)
              end do
            end do
          end do
          !$OMP END PARALLEL DO
        case (7)
          !$OMP PARALLEL DO COLLAPSE(3)
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                dghdx(i,j,k) = center_diff6(ght(i-3:i+3,j,k), dx)
                dghdy(i,j,k) = center_diff6(ght(i,j-3:j+3,k), dy)
              end do
            end do
          end do
          !$OMP END PARALLEL DO
        end select
        end associate
      end if
    end do
    end associate

  end subroutine calc_grad_gh

end module operators_adv_mod
